/*
 *   Copyright 2012, Thomas Kerber
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 */
package milk.jpatch;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.TreeMap;

import static milk.jpatch.Util.logger;
import milk.jpatch.code.CodePointHook;

import org.apache.bcel.classfile.*;
import org.apache.bcel.generic.InstructionHandle;
import org.apache.bcel.generic.ConstantPoolGen;
import org.apache.bcel.generic.InstructionList;

/**
 * Maps references to the old cpool (of the modified file) to ones in the new
 * cpool.
 * 
 * This is also in charge of shifting cpool references of existing objects to
 * the new cpool.
 * 
 * Note that this map is generated at patch time, -not- at diff time.
 * 
 * @author Thomas Kerber
 * @version 1.0.2
 */
public class CPoolMap extends TreeMap<Integer, Integer>{
    
    private static final long serialVersionUID = -8146832398911942630L;
    
    /**
     * The attribute names recognized by this class.
     * 
     * Yes, this is a "final" mutable object. Just <em>treat</em> it as final.
     */
    public static final Set<String> RECOGNIZED_ATTRIBS =
            new HashSet<String>(Arrays.asList(new String[]{
                    "ConstantValue",
                    "Code",
                    "Exceptions",
                    "InnerClasses",
                    "EnclosingMethod",
                    "Synthetic",
                    "Signature",
                    "SourceFile",
                    "SourceDebugExtension",
                    "LineNumberTable",
                    "LocalVariableTable",
                    "LocalVariableTypeTable",
                    "Deprecated",
                    "RuntimeVisibleAnnotations",
                    "RuntimeInvisibleAnnotations",
                    "RuntimeVisibleParameterAnnotations",
                    "RuntimeInvisibleParameterAnnotations",
                    "AnnotationDefault",
                    "StackMapTable"
            }));
    
    /**
     * The constant pool from which substitutions are made.
     */
    public ConstantPool from = null;
    /**
     * The constant pool to which substitutions are made.
     */
    public ConstantPool to = null;
    /**
     * The ConstantPoolGen form of "from".
     */
    public ConstantPoolGen fromGen = null;
    /**
     * The ConstantPoolGen form of "to".
     */
    public ConstantPoolGen toGen = null;
    
    /**
     * Generates a CPoolMap.
     * @param to The constant pool whose indexes don't change and which counts
     *         as the "destination" cpool.
     * @param from The constant pool from which is substituted.
     * @return The CPoolMap.
     */
    public static CPoolMap generate(ConstantPool to, ConstantPool from){
        CPoolMap map = new CPoolMap();
        map.toGen = new ConstantPoolGen(to);
        map.from = from;
        map.fromGen = new ConstantPoolGen(map.from);
        Constant[] consts = map.fromGen.getConstantPool().getConstantPool();
        // index 0 is always null.
        for(int i = 1; i < consts.length; i++){
            if(consts[i] == null)
                continue;
            int newPos = map.toGen.addConstant(consts[i], map.fromGen);
            map.put(i, newPos);
        }
        map.to = map.toGen.getFinalConstantPool();
        
        return map;
    }
    
    /**
     * Shifts the constant pool from "from" to "to".
     * @param f The object to shift.
     * @return The shifted object.
     */
    public FieldOrMethod applyTo(FieldOrMethod f){
        
        if(f == null)
            return null;
        
        // Delegates the attribute substitution.
        Attribute[] oldAttrs = f.getAttributes();
        Attribute[] newAttrs = new Attribute[oldAttrs.length];
        for(int i = 0; i < oldAttrs.length; i++)
            newAttrs[i] = this.applyTo(oldAttrs[i]);
        
        // Substitutes name index, signature index and uses the new cpool.
        if(f instanceof Field)
            return new Field(f.getAccessFlags(),
                    this.get(f.getNameIndex()),
                    this.get(f.getSignatureIndex()),
                    newAttrs,
                    this.to);
        else // Method
            return new Method(f.getAccessFlags(),
                    this.get(f.getNameIndex()),
                    this.get(f.getSignatureIndex()),
                    newAttrs,
                    this.to);
    }

    /**
     * Shifts the constant pool from "from" to "to".
     * @param f The object to shift.
     * @return The shifted object.
     */
    public Field applyTo(Field f){
        return (Field)this.applyTo((FieldOrMethod)f);
    }

    /**
     * Shifts the constant pool from "from" to "to".
     * @param m The object to shift.
     * @return The shifted object.
     */
    public Method applyTo(Method m){
        return (Method)this.applyTo((FieldOrMethod)m);
    }

    /**
     * Shifts the constant pool from "from" to "to".
     * @param c The object to shift.
     * @return The shifted object.
     */
    public InnerClass applyTo(InnerClass c){
        
        if(c == null)
            return null;
        
        return new InnerClass(this.get(c.getInnerClassIndex()),
                this.get(c.getOuterClassIndex()),
                this.get(c.getInnerNameIndex()),
                c.getInnerAccessFlags());
    }

    /**
     * Shifts the constant pool from "from" to "to".
     * @param v The object to shift.
     * @return The shifted object.
     */
    public LocalVariable applyTo(LocalVariable v){
        
        if(v == null)
            return null;
        
        return new LocalVariable(v.getStartPC(),
                v.getLength(),
                this.get(v.getNameIndex()),
                this.get(v.getSignatureIndex()),
                v.getIndex(),
                this.to);
    }

    /**
     * Shifts the constant pool from "from" to "to".
     * @param v The object to shift.
     * @return The shifted object.
     */
    public ElementValue applyTo(ElementValue v){
        
        if(v == null)
            return null;
        
        if(v instanceof SimpleElementValue){
            return new SimpleElementValue(
                    v.getElementValueType(),
                    this.get(((SimpleElementValue)v).getIndex()),
                    this.to);
        }
        
        else if(v instanceof EnumElementValue){
            return new EnumElementValue(
                    v.getElementValueType(),
                    this.get(((EnumElementValue)v).getTypeIndex()),
                    this.get(((EnumElementValue)v).getValueIndex()),
                    this.to);
        }
        
        else if(v instanceof ClassElementValue){
            return new ClassElementValue(
                    v.getElementValueType(),
                    this.get(((ClassElementValue)v).getIndex()),
                    this.to);
        }
        
        else if(v instanceof AnnotationElementValue){
            return new AnnotationElementValue(
                    v.getElementValueType(),
                    this.applyTo(((AnnotationElementValue)v).
                            getAnnotationEntry()),
                    this.to);
        }
        
        else if(v instanceof ArrayElementValue){
            ElementValue[] oldVals = ((ArrayElementValue)v).
                    getElementValuesArray();
            ElementValue[] newVals = new ElementValue[oldVals.length];
            for(int i = 0; i < oldVals.length; i++)
                newVals[i] = this.applyTo(oldVals[i]);
            
            return new ArrayElementValue(
                    v.getElementValueType(),
                    newVals,
                    this.to);
        }
        
        throw new IllegalStateException(
                "Unreachable position reached. (Unknown ElementValue type)");
    }

    /**
     * Shifts the constant pool from "from" to "to".
     * @param p The object to shift.
     * @return The shifted object.
     */
    public ElementValuePair applyTo(ElementValuePair p){
        if(p == null)
            return null;
        
        return new ElementValuePair(
                this.get(p.getNameIndex()),
                this.applyTo(p.getValue()),
                this.to);
    }

    /**
     * Shifts the constant pool from "from" to "to".
     * @param e The object to shift.
     * @return The shifted object.
     */
    public AnnotationEntry applyTo(AnnotationEntry e){
        if(e == null)
            return null;
        
        AnnotationEntry ret = new AnnotationEntry(
                this.get(e.getAnnotationTypeIndex()),
                this.to,
                e.isRuntimeVisible());
        
        for(ElementValuePair p : e.getElementValuePairs())
            ret.addElementNameValuePair(this.applyTo(p));
        
        return ret;
    }

    /**
     * Shifts the constant pool from "from" to "to".
     * @param c The object to shift.
     * @return The shifted object.
     */
    public Code applyTo(Code c){
        return applyTo(c, new ArrayList<CodePointHook>());
    }
    
    /**
     * Shifts the constant pool from "from" to "to".
     * @param il The object to shift.
     * @return The shifted object.
     */
    public InstructionList applyTo(InstructionList il){
        il.replaceConstantPool(fromGen, toGen);
        il.setPositions();
        return il;
    }
    
    /**
     * Shifts the constant pool from "from" to "to".
     * @param c The object to shift.
     * @param cphs The CodePointHooks to also shift accordingly.
     * @return The shifted object.
     */
    public Code applyTo(Code c, List<CodePointHook> cphs){
        /*
         * Normally, the handles -should- stay the same, as the code isn't
         * actually being changed. However some minor alterations could happen,
         * which in themselves would already be catastrophic for all and any
         * handles. Therefore, they get stored and updated.
         */
        
        // Handles are stored in triples, representing the start pc, end pc
        // and handler pc of exception table in that order.
        // Code point hooks are stored after these.
        if(c == null)
            return null;
        
        CodeException[] oldExceptions = c.getExceptionTable();
        
        InstructionHandle[] handles =
                new InstructionHandle[oldExceptions.length * 3 + cphs.size()];
        
        InstructionList il = new InstructionList(c.getCode());
        
        for(int i = 0; i < oldExceptions.length; i++){
            handles[3 * i] = il.findHandle(oldExceptions[i].getStartPC());
            handles[3 * i + 1] = il.findHandle(oldExceptions[i].getEndPC());
            handles[3 * i + 2] = il.findHandle(oldExceptions[i].
                    getHandlerPC());
        }
        for(int i = 0; i < cphs.size(); i++){
            handles[3 * oldExceptions.length + i] =
                    il.findHandle(cphs.get(i).codePoint);
        }
        
        il.replaceConstantPool(this.fromGen, this.toGen);
        il.setPositions();
        
        CodeException[] newExceptions =
                new CodeException[oldExceptions.length];
        for(int i = 0; i < oldExceptions.length; i++){
            if(oldExceptions[i] == null){
                newExceptions[i] = null;
                continue;
            }
            
            int catchType = oldExceptions[i].getCatchType();
            newExceptions[i] = new CodeException(
                    handles[i * 3].getPosition(),
                    handles[i * 3 + 1].getPosition(),
                    handles[i * 3 + 2].getPosition(),
                    catchType == 0 ? 0 : this.get(catchType));
        }
        
        for(int i = 0; i < cphs.size(); i++){
            cphs.get(i).codePoint =
                    handles[3 * oldExceptions.length + i].getPosition();
        }
        
        Attribute[] oldAttribs = c.getAttributes();
        Attribute[] newAttribs = new Attribute[oldAttribs.length];
        for(int i = 0; i < oldAttribs.length; i++)
            newAttribs[i] = this.applyTo(oldAttribs[i]);
        
        return new Code(
                this.get(c.getNameIndex()),
                c.getLength(),
                c.getMaxStack(),
                c.getMaxLocals(),
                il.getByteCode(),
                newExceptions,
                newAttribs,
                this.to);
    }
    
    /**
     * Shifts the constant pool from "from" to "to".
     * @param a The object to shift.
     * @return The shifted object.
     */
    public Attribute applyTo(Attribute a){
        if(a == null)
            return null;
        
        int newName = this.get(a.getNameIndex());
        
        // Procedure defined for each of the predefined Attributes as per JVM
        // Spec. §4.8 (All cpool references are shifted)
        
        if(a instanceof SourceFile){
            return new SourceFile(newName,
                    a.getLength(),
                    this.get(((SourceFile)a).getSourceFileIndex()),
                    this.to);
        }
        
        else if(a instanceof ConstantValue){
            return new ConstantValue(newName,
                    a.getLength(),
                    this.get(((ConstantValue)a).getConstantValueIndex()),
                    this.to);
        }
        
        else if(a instanceof Code){
            // Code is kinda a biggie, so it's been extracted. (Plus if you
            // want to shift just a code, you generally don't want an Attribute
            // returning.
            return this.applyTo((Code)a);
        }
        
        else if(a instanceof ExceptionTable){
            // Maps all exception indexes
            int[] oldIndexTable = ((ExceptionTable)a).getExceptionIndexTable();
            int[] newIndexTable = new int[oldIndexTable.length];
            for(int i = 0; i < oldIndexTable.length; i++)
                newIndexTable[i] = this.get(oldIndexTable[i]);
            
            return new ExceptionTable(newName,
                    a.getLength(),
                    newIndexTable,
                    this.to);
        }
        
        else if(a instanceof InnerClasses){
            // Delegates the inner class mappings.
            InnerClass[] oldICs = ((InnerClasses)a).getInnerClasses();
            InnerClass[] newICs = new InnerClass[oldICs.length];
            for(int i = 0; i < oldICs.length; i++)
                newICs[i] = this.applyTo(oldICs[i]);
            
            return new InnerClasses(newName,
                    a.getLength(),
                    newICs,
                    this.to);
        }
        
        // Synthetic skipped due to no cpool refs.
        
        // LineNumberTable skipped due to no cpool refs.
        
        else if(a instanceof LocalVariableTable){
            // Delegates the local variable mappings.
            LocalVariable[] oldLVs = ((LocalVariableTable)a).
                    getLocalVariableTable();
            LocalVariable[] newLVs = new LocalVariable[oldLVs.length];
            for(int i = 0; i < oldLVs.length; i++)
                newLVs[i] = this.applyTo(oldLVs[i]);
            
            return new LocalVariableTable(newName,
                    a.getLength(),
                    newLVs,
                    this.to);
        }
        
        else if(a instanceof EnclosingMethod){
            EnclosingMethod e = (EnclosingMethod)a.copy(this.to);
            e.setEnclosingClassIndex(this.get(e.getEnclosingClassIndex()));
            e.setEnclosingMethodIndex(this.get(e.getEnclosingMethodIndex()));
            e.setNameIndex(newName);
            return e;
        }
        
        else if(a instanceof Signature){
            return new Signature(newName,
                    a.getLength(),
                    this.get(((Signature)a).getSignatureIndex()),
                    this.to);
        }
        
        else if(a instanceof LocalVariableTypeTable){
            // Delegates the local variable mappings.
            LocalVariable[] oldLVs = ((LocalVariableTypeTable)a).
                    getLocalVariableTypeTable();
            LocalVariable[] newLVs = new LocalVariable[oldLVs.length];
            for(int i = 0; i < oldLVs.length; i++)
                newLVs[i] = this.applyTo(oldLVs[i]);
            
            return new LocalVariableTypeTable(newName,
                    a.getLength(),
                    newLVs,
                    this.to);
            
        }
        
        // SourceDebugExtension skipped, as it is a) not in BCEL and b) has no
        // (known) cpool refs.
        
        // Deprecated skipped due to no cpool refs.
        
        else if(a instanceof RuntimeVisibleAnnotations){
            AnnotationEntry[] oldAnnots = ((Annotations)a).
                    getAnnotationEntries();
            AnnotationEntry[] newAnnots =
                    new AnnotationEntry[oldAnnots.length];
            for(int i = 0; i < oldAnnots.length; i++)
                newAnnots[i] = this.applyTo(oldAnnots[i]);
            
            RuntimeVisibleAnnotations cp =
                    (RuntimeVisibleAnnotations)a.copy(this.to);
            cp.setAnnotationTable(newAnnots);
            cp.setNameIndex(newName);
            return cp;
        }
        
        else if(a instanceof RuntimeVisibleParameterAnnotations){
            ParameterAnnotationEntry[] oldAnnots = ((ParameterAnnotations)a).
                    getParameterAnnotationEntries();
            ParameterAnnotationEntry[] newAnnots =
                    new ParameterAnnotationEntry[oldAnnots.length];
            for(int i = 0; i < oldAnnots.length; i++){
                AnnotationEntry[] oldEntries =
                        oldAnnots[i].getAnnotationEntries();
                AnnotationEntry[] newEntries =
                        new AnnotationEntry[oldEntries.length];
                for(int f = 0; f < oldEntries.length; f++)
                    newEntries[f] = this.applyTo(oldEntries[f]);
                newAnnots[i] = new ParameterAnnotationEntry(newEntries);
            }
            
            RuntimeVisibleParameterAnnotations pa =
                    (RuntimeVisibleParameterAnnotations)a.copy(
                            this.to);
            pa.setParameterAnnotationTable(newAnnots);
            pa.setNameIndex(newName);
            return pa;
        }
        
        else if(a instanceof RuntimeInvisibleAnnotations){
            AnnotationEntry[] oldAnnots = ((Annotations)a).
                    getAnnotationEntries();
            AnnotationEntry[] newAnnots =
                    new AnnotationEntry[oldAnnots.length];
            for(int i = 0; i < oldAnnots.length; i++)
                newAnnots[i] = this.applyTo(oldAnnots[i]);
            
            RuntimeInvisibleAnnotations cp =
                    (RuntimeInvisibleAnnotations)a.copy(this.to);
            cp.setAnnotationTable(newAnnots);
            cp.setNameIndex(newName);
            return cp;
            
        }
        
        else if(a instanceof RuntimeInvisibleParameterAnnotations){
            ParameterAnnotationEntry[] oldAnnots = ((ParameterAnnotations)a).
                    getParameterAnnotationEntries();
            ParameterAnnotationEntry[] newAnnots =
                    new ParameterAnnotationEntry[oldAnnots.length];
            for(int i = 0; i < oldAnnots.length; i++){
                AnnotationEntry[] oldEntries =
                        oldAnnots[i].getAnnotationEntries();
                AnnotationEntry[] newEntries =
                        new AnnotationEntry[oldEntries.length];
                for(int f = 0; f < oldEntries.length; f++)
                    newEntries[f] = this.applyTo(oldEntries[f]);
                newAnnots[i] = new ParameterAnnotationEntry(newEntries);
            }
            
            RuntimeInvisibleParameterAnnotations pa =
                    (RuntimeInvisibleParameterAnnotations)a.copy(
                            this.to);
            pa.setParameterAnnotationTable(newAnnots);
            pa.setNameIndex(newName);
            return pa;
        }
        
        else if(a instanceof AnnotationDefault){
            return new AnnotationDefault(
                    newName,
                    a.getLength(),
                    this.applyTo(((AnnotationDefault) a).getDefaultValue()),
                    this.to);
        }
        
        // StackMapTable skipped, as it gets auto-removed anyway.
        
        // TODO: BootstrapMethods
        
        if(!RECOGNIZED_ATTRIBS.contains(a.getName()))
            // Unknown attribute!
            logger.warning("Unknown Attribute encountered during cpool " +
                    "shift. (" + a.getName() + ") The attribute may contain " +
                    "invalid or false cpool references.");
        Attribute cp = a.copy(this.to);
        cp.setNameIndex(newName);
        return cp;
    }
    
}
